#include "StdAfx.h"
#include "WebConn.h"

#include <iostream>
#include <string>

#include "roles/client.hpp"
#include "websocketpp.hpp"

class chat_client_handler : public websocketpp::client::handler
{
public:
	chat_client_handler(CWebSocketConn* pWsConn, WSONMSGCALLBACK OnMsgCallBack);
	virtual ~chat_client_handler();
	void on_fail(connection_ptr con);
	void on_open(connection_ptr con);
	void on_close(connection_ptr con);
	void on_message(connection_ptr con, message_ptr msg);

	BOOL send(const std::string &msg);
	void close();

public:
	connection_ptr m_con;
	CWebSocketConn* m_pWsConn;
};

class CWebSocketConn
{
public:
	CWebSocketConn(CWebWsConn* pParent, LPCTSTR lpszUri, WSONMSGCALLBACK OnMsgCallBackFunc)
		: m_cs()
		, m_pParent(pParent)
		, m_hWait(CreateEvent(NULL, FALSE, FALSE, NULL))
		, m_ChatClientHandler(new chat_client_handler(this, OnMsgCallBackFunc))
		, m_endpoint(m_ChatClientHandler)
		, m_pThread(NULL)
	{
		USES_CONVERSION;

		m_endpoint.alog().unset_level(websocketpp::log::alevel::ALL);
		m_endpoint.elog().unset_level(websocketpp::log::elevel::ALL);
/*
		m_endpoint.elog().set_level(websocketpp::log::elevel::RERROR);
		m_endpoint.elog().set_level(websocketpp::log::elevel::FATAL);
*/
		m_con = m_endpoint.get_connection(T2A(lpszUri));

		m_con->add_request_header("User-Agent","WebSocket++/0.2.0 WebSocket++Chat/0.2.0");
		m_con->add_subprotocol("com.zaphoyd.websocketpp.chat");

		m_con->set_origin("http://zaphoyd.com");

		m_endpoint.connect(m_con);

		m_pThread = new boost::thread(boost::bind(&websocketpp::client::run, &m_endpoint, false));

		if (WaitForSingleObject(m_hWait, 30*1000) == WAIT_OBJECT_0)
		{
			if (!m_ChatClientHandler->m_con)
			{
				Release();
				throw 0;
			}

			m_bIsValid = TRUE;
		}
		else
		{
			Release();
			throw 0;
		}
	}

	~CWebSocketConn()
	{
		if (m_bIsValid)
		{
			Release();
		}
	}

	void Release()
	{
		CSingleLock sl(&m_cs);

		m_bIsValid = FALSE;
		sl.Lock();
		if (m_hWait != NULL)
		{
			CloseHandle(m_hWait);
			m_hWait = NULL;
		}
		if (m_pThread != NULL)
		{
			m_endpoint.close_all();
			(*m_con).close(websocketpp::close::status::GOING_AWAY, "");
			m_ChatClientHandler->close();
			
			if (m_pThread->joinable() && boost::this_thread::get_id() != m_pThread->get_id())
			{
				m_pThread->join();
			}
			delete m_pThread;
			m_pThread = NULL;
		}
		sl.Unlock();
	}

	BOOL Send(LPCTSTR lpszMsg)
	{
		USES_CONVERSION;

		return m_ChatClientHandler->send(T2A(lpszMsg));
	}

	static void OnErrorCallBack(CWebSocketConn* pThis)
	{
		pThis->Release();
		pThis->m_pParent->WsReconnect();
	}

public:
	CWebWsConn* m_pParent;
	CCriticalSection m_cs;
	HANDLE m_hWait;
	WSONMSGCALLBACK m_OnMsgCallBackFunc;
	boost::shared_ptr<chat_client_handler> m_ChatClientHandler;
	websocketpp::client m_endpoint;
	websocketpp::client::connection_ptr m_con;
	boost::thread *m_pThread;
	BOOL m_bIsValid;
};


CWebWsConn::CWebWsConn(void)
: m_pWebSocketConn(NULL)
, m_bFirstTry(TRUE)
, m_hHeartBeatThread(NULL)
, m_bReconnecting(FALSE)
{
}

CWebWsConn::~CWebWsConn(void)
{
	WsDisconnect();
}

BOOL CWebWsConn::WsConnect(LPCTSTR lpszWsUrl, WSONMSGCALLBACK pOnMsgCallBackFunc, WSONRECONNECTED pOnReconnectedCallBackFunc)
{
	if (pOnReconnectedCallBackFunc != NULL)
	{
		m_pOnReconnectedCallBackFunc = pOnReconnectedCallBackFunc;
	}

	WsDisconnect();
	try
	{
		m_csWsUrl = lpszWsUrl;
		m_pOnMsgCallBackFunc = pOnMsgCallBackFunc;
		m_pWebSocketConn = new CWebSocketConn(this, lpszWsUrl, pOnMsgCallBackFunc);
	}
	catch (...)
	{
		if (m_pWebSocketConn != NULL)
		{
			delete m_pWebSocketConn;
			m_pWebSocketConn = NULL;
		}
		return FALSE;
	}

	m_bFirstTry = FALSE;
	WsStartHeartBeat();

	return TRUE;
}

void CWebWsConn::WsDisconnect()
{
	WsStopHeartBeat();
	if (m_pWebSocketConn != NULL)
	{
		delete m_pWebSocketConn;
		m_pWebSocketConn = NULL;
	}
}

DWORD CALLBACK ThreadDoReconnect(LPVOID lpParam)
{
	CWebWsConn* pConn = (CWebWsConn*)lpParam;
	DWORD dwRetryCounter = 0;

	while (TRUE)
	{
		if (pConn->WsConnect(pConn->m_csWsUrl, pConn->m_pOnMsgCallBackFunc, NULL))
		{
			if (pConn->m_pOnReconnectedCallBackFunc != NULL)
			{
				pConn->m_pOnReconnectedCallBackFunc(pConn);
			}
			pConn->m_bReconnecting = FALSE;
			break;
		}

		dwRetryCounter++;

		TRACE(_T("I\'m trying reconnect %d times..."), dwRetryCounter);

		if (dwRetryCounter <= 8)
		{
			Sleep((DWORD)pow(2.0, (int)dwRetryCounter) * 1000);
		}
		else
		{
			Sleep(5*60*1000);
		}
	}
	return 0;
}

void CWebWsConn::WsReconnect()
{
	if (!m_bFirstTry && !m_bReconnecting)
	{
		m_bReconnecting = TRUE;
		CloseHandle(CreateThread(NULL, 0, ThreadDoReconnect, this, 0, NULL));
	}
}

BOOL CWebWsConn::WsSend(LPCTSTR lpszMsg)
{
	if (m_pWebSocketConn == NULL)
	{
		return FALSE;
	}
	if (m_pWebSocketConn->m_bIsValid)
	{
		return m_pWebSocketConn->Send(lpszMsg);
	}
	else
	{
		WsReconnect();
		return FALSE;
	}
}

DWORD CALLBACK ThreadHeartBeat(LPVOID lpParam)
{
	CWebWsConn* pConn = (CWebWsConn*)lpParam;

//	Sleep(HEARTBEAT_INTERVAL);

	while (!pConn->m_bStopHeartBeat && pConn->m_pWebSocketConn != NULL && pConn->m_pWebSocketConn->m_bIsValid && pConn->m_pWebSocketConn->m_ChatClientHandler.get() && pConn->m_pWebSocketConn->m_ChatClientHandler->m_con)
	{
		pConn->m_pWebSocketConn->m_ChatClientHandler->m_con->ping(std::string(""));
		Sleep(HEARTBEAT_INTERVAL);
	}

	if (!pConn->m_bStopHeartBeat)
	{
		pConn->WsReconnect();
	}

	return 0;
}


void CWebWsConn::WsStartHeartBeat()
{
	WsStopHeartBeat();
	m_bStopHeartBeat = FALSE;
	m_hHeartBeatThread = CreateThread(NULL, 0, ThreadHeartBeat, this, 0, NULL);
}

void CWebWsConn::WsStopHeartBeat()
{
	m_bStopHeartBeat = TRUE;
	if (m_hHeartBeatThread != NULL)
	{
		WaitForSingleObject(m_hHeartBeatThread, HEARTBEAT_INTERVAL);
		CloseHandle(m_hHeartBeatThread);
		m_hHeartBeatThread = NULL;
	}
}

chat_client_handler::chat_client_handler(CWebSocketConn* pWsConn, WSONMSGCALLBACK OnMsgCallBack)
{
	m_pWsConn = pWsConn;
	m_pWsConn->m_OnMsgCallBackFunc = OnMsgCallBack;
}

chat_client_handler::~chat_client_handler()
{
	close();
}

void chat_client_handler::on_fail(connection_ptr con)
{
	SetEvent(m_pWsConn->m_hWait);

	CWebSocketConn::OnErrorCallBack(m_pWsConn);
}

void chat_client_handler::on_open(connection_ptr con)
{
	m_con = con;
	SetEvent(m_pWsConn->m_hWait);
}

void chat_client_handler::on_close(connection_ptr con)
{
	m_con = connection_ptr();
}

void chat_client_handler::on_message(connection_ptr con, message_ptr msg)
{
	USES_CONVERSION;
	m_pWsConn->m_OnMsgCallBackFunc(A2T(msg->get_payload().c_str()));
}

BOOL chat_client_handler::send(const std::string &msg)
{
	if (!m_con)
	{
		return FALSE;
	}

	m_con->send(msg);

	return TRUE;
}

void chat_client_handler::close()
{
	if (!m_con)
	{
		return;
	}

	m_con->close(websocketpp::close::status::GOING_AWAY,"");
}